#!/usr/bin/env python
#
# Author: LHC simple program
# simple example class function object 

#class Dog:
#   def __init__(self, name):
#       self.name = name
#   
#   def bark(self):
#       print(f"{self.name} is barking!")

#my_dog = Dog("Buddy")  # 创建一个Dog类的对象
#my_dog.bark()  	# 调用对象的bark方法

'''
A simple example to run RHF 

'''

from pyscf import scf,gto,lib
import pyscf
import numpy as np
import scipy
import opt
import hessian

#ol = pyscf.M(
#   atom = 'H 0 0 0; F 0 0 1.1',  # in Angstrom
#   basis = 'sto-3g',
#   symmetry = True,
#
mol = gto.Mole() #这行代码初始化了一个gto.Mole对象，gto是PySCF中用于构建分子的模块。Mole是该模块中的类，用于存储和处理分子相关的信息（如原子类型、位置、基组等）。
mol.atom = """
N  0.            0.   0.
H  0.            1.  -0.2
H  0.8660254038 -0.5 -0.2
H -0.8660254038 -0.5 -0.2
""" 		     ## unit:Angstrom
mol.basis = "6-31G"  ## 22 orbital
mol.verbose = 0
mol.build()
#pyscf interface
A_t=mol.atom_coords() ##
T = mol.intor("int1e_kin")
V = mol.intor('int1e_nuc')
S = mol.intor('int1e_ovlp')
X = scipy.linalg.fractional_matrix_power(S,-0.5)##对S矩阵求S^(-1/2)
eri=mol.intor('int2e')
tuv=mol.intor("int1e_r") ##Calculate the dipole moment  3D(t,u,v)
#print(eri.shape)
Hcore=T + V

natm=mol.natm    #atom number
nmo =nao=mol.nao #atom orbital number
nocc=mol.nelec[0] #[0]:Alpha ele [1]:Beta ele
#print(nocc)
so=slice(0,nocc) #左闭右开[0,9)
#print(so)
A_t = mol.atom_coords()
print(A_t)
r_ABt =A_t[:,None,:] - A_t[None,:,:]
r_AB = np.linalg.norm(r_ABt,axis=-1)
r_AB += np.diag(np.ones(mol.natm)*np.inf)
Z_A = mol.atom_charges()
#AA_t=Z_A[:,None,None]
#print(AA_t)
#A_charge=mol.atom_charges()[:,None]*mol.atom_charges()
A_charge=Z_A[:,None]*Z_A
#print(A_charge)
E_nuc=0.5 *(A_charge / r_AB).sum()
print('Neucleus energy   ',E_nuc,'a.u.')
###Provide a first guess using pyscf
mf=scf.RHF(mol)
#D=mf.get_init_guess(mol,key='minao') # 1)
#D=mf.get_init_guess(mol,key='huckel')# 2)
D=mf.get_init_guess(mol,key='atom')   # 3)
#D=mf.get_init_guess(mol,key='1e')    # 4)
#print(D.shape)
###
D_old=np.zeros((nao,nao))
E_old=np.inf
E_ele=0         #初始能量
energy_tol=1e-6 #收敛阈值
density_tol=1e-6
#print(E_old)
count=0
while not(np.allclose(D, D_old, atol=density_tol) and abs(E_ele - E_old) < energy_tol):
   #print('1')
    if count > 500:
        raise ValueError("SCF not converged!")##大于最大迭代步数时报错（TRUE）
    count += 1
    D_old=D
    E_old=E_ele
    F = Hcore + np.einsum("uvkl, kl -> uv",eri,D) - 0.5*np.einsum("ukvl, kl -> uv",eri,D)
    Fp= X.T @ F @ X ##矩阵乘法(Fortran:matmul)
    #e, C = scipy.linalg.eigh(F, S)  # Solve FC = SCe
    e,Cp = np.linalg.eigh(Fp) ##求解特征值和特征向量
    C=X @ Cp
    D= 2 * C[:,so] @  C[:,so].T
    #E_ele=0.5*np.sum(D.T*(Hcore+F))
    E_ele = (Hcore * D).sum() + 0.5 * np.einsum("uvkl, uv, kl ->", eri, D, D) - 0.25 * np.einsum("ukvl, uv, kl ->", eri, D, D)
E_tot=E_ele + E_nuc
#print(e.shape)
print("SCF Converged in  ", count, " loops")
print("Electronic energy ", E_ele, " a.u.")
print("Total energy      ", E_tot, " a.u.")
for i in range(nao):
    #print(f"MO{i}        ",e[i],   " a.u.")
    print(f"MO {i}        {e[i]:.6f} a.u.",)
print("----------------- ")
#print(lib.param.BOHR)
#stop
###calculate hessian
hess=hessian.cal_hess(A_t,mol.natm,mol.basis,mol.atom)
print(hess)
###end calculate
### opt geom
coord1=opt.opt_geom(A_t,mol.atom,mol.basis)
print(coord1*(lib.param.BOHR))
###end opt geom
###Calculate the dipole moment###
dipole_m=-np.einsum("tuv,uv->t",tuv,D)+np.einsum("a,at->t",Z_A,A_t)
print(dipole_m)
##occ
#co=C[:,so]
#print(co.shape)
###
                    ###MP2###
nocc,nmo,nao,nbas=mol.nelec[0],mol.nao,mol.nao,mol.nbas
nvir=nmo-nocc
so,sv,sa=slice(0,nocc),slice(nocc,nmo),slice(0,nbas)
#print(so,sv,sa)
Co,Cv=C[:,so],C[:,sv]
eo,ev=e[so],e[sv]
print('Co and Cv shape:',Co.shape,Cv.shape)
print('eo and ev shape:',eo.shape,ev.shape)
eri0_iajb = np.einsum("ui, va, uvkl, kj, lb -> iajb", Co, Cv, eri, Co, Cv)
print(eri0_iajb.shape)
#D_ijab=np.zeros((nocc, nvir, nocc, nvir))
# for i in nocc:
#     for j in nvir:
#         for k in nocc:
#             for l in nvir:
#                 D_ijab(i,j,k,l)=eo(i)-ev(j)+eo(k)-ev(l)   
D_iajb = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]
print(D_iajb.shape)
t_iajb = eri0_iajb / D_iajb
T_iajb = 2 * t_iajb - t_iajb.swapaxes(1, 3)
E_mp2=np.einsum("iajb, iajb, iajb ->", T_iajb, t_iajb, D_iajb)
print("MP2 energy ", E_mp2, " a.u.")
print("EMP2 energy ",E_tot+E_ele, " a.u.")









